module evolve

  use dimensions

  complex, private      :: ExpKv(jz, jv)
  complex, private      :: ExpKx(jx, jDirection)
  complex, private      :: ExpKy(jy, jDirection)
  complex, private      :: ExpKz(jz)

  complex, private      :: uxMod(jx)
  complex, private      :: uyMod(jy)

  real, private         ::  weight 

contains

!-------------------------------------------------------------------------------------
subroutine evolve_init()
!-------------------------------------------------------------------------------------

  use parameters, only : dt, kz, vPhase, ux0, uy0, Dxxxx, &
                         iWaveFrame, iunit, Lx, Ly

  use fft_wrap, only   : fft_forward, fft_backward

  implicit none

  integer              :: iDirection, ikz
  real                 :: v0

  !-- setting modulation vectors --

  uxMod(:) = exp( - 14 * xvec(:)**6 / (Lx/2)**6 )
  uyMod(:) = exp( - 14 * yvec(:)**6 / (Ly/2)**6 )

  if (Lx .eq. 0) uxMod(:) = 1.
  if (Ly .eq. 0) uyMod(:) = 1.
 
  call fft_forward(jx,uxMod,uxMod)
  uxMod=uxMod/float(jx)
  call deAlias(jx,uxMod)        
  call fft_backward(jx,uxMod,uxMod)

  call fft_forward(jy,uyMod,uyMod)
  uyMod=uyMod/float(jy)
  call deAlias(jy,uyMod)        
  call fft_backward(jy,uyMod,uyMod)



  !-- initialization --

  ! weight = 0.5 since each direction has 1/2 density,
  ! goes to 1/4 is 3D (1/6 with hexagonal symmetry)

  weight = 1./jDirection

  v0 = -vPhase

  ExpKz(:) = cexp( iunit * kZvec(:) * v0 * dt )
      
  do ikz=1,nkz
     ExpKv(ikz,:) = cexp( -iunit*vvec(:)*kZvec(ikz)*dt - Dxxxx * kZvec(ikz)**4 *dt    )
  enddo

  do iDirection = 1,jDirection
     ExpKx(:,iDirection) = exp( - iunit * kXvec(:)*ux0(iDirection) * dt  )
     ExpKy(:,iDirection) = exp( - iunit * kYvec(:)*uy0(iDirection) * dt  )
  enddo

end subroutine evolve_init


!-------------------------------------------------------------------------------------
subroutine evolve_F_advect(dfu, &                              ! input/output
                           uxMat, uyMat, f0vec)                ! input
!-------------------------------------------------------------------------------------

  use fft_wrap, only   : fft_forward, fft_backward
  use parameters, only : Dxxxx, dt, iBcPerpPeriodic

  implicit none

  real,    intent(inout)   :: dfu(jv, jz, jy, jx, jDirection)
  real,    intent(in)      :: uxMat(jz, jy, jx, jDirection)
  real,    intent(in)      :: uyMat(jz, jy, jx, jDirection)
  real,    intent(in)      :: f0vec(jv)

  complex                  :: fz(jz)
  real                     :: axVar(jx), fxVec(jx)
  real                     :: ayVar(jy), fyVec(jy)
  integer                  :: ix, iy, iz, iv, iDirection


 !------------------------------------------

    
    
  do iDirection = 1,jDirection

   
     !-- advect in z: df/dt = -v df/dz --

     do ix=1,jx
        do iy=1, jy
           do iv=1,jv
              fz(:) = dfu(iv,:,iy,ix,iDirection)
              call fft_forward(jz,fz,fz)
              fz = fz/float(jz)
              call dealias(jz,fz)
              fz(:) = fz(:)*ExpKv(:,iv)
              call fft_backward(jz,fz,fz)
              dfu(iv,:,iy,ix,iDirection) = fz(:)
           enddo
        end do
     enddo
    

     !-- advect in x: d(f+f0)/dt = -d/dx [(f+f0) Ux(x,y)] --

     do iy=1,jy
        do iz=1,jz                      
           axVar(:) = uxMat(iz,iy,:,iDirection)           
           do iv=1,jv
              fxVec(:) = dfu(iv,iz,iy,:,iDirection) + weight * f0vec(iv)
              call UpWind2ndOrder_Variable_a( fxVec,  dt, dx, axVar, jx, &
                   weight*f0vec(iv),  weight*f0vec(iv), iBcPerpPeriodic )
              dfu(iv,iz,iy,:,iDirection) = fxVec(:)  - weight * f0vec(iv) 
           enddo
        enddo
     end do

     !-- advect in y: d(f+f0)/dt = -d/dy [(f+f0) Uy(x,y)] --

     do ix=1,jx
        do iz=1,jz                      
           ayVar(:) = uyMat(iz,:,ix,iDirection)           
           do iv=1,jv
              fyVec(:) = dfu(iv,iz,:,ix,iDirection) + weight * f0vec(iv)
              call UpWind2ndOrder_Variable_a( fyVec,  dt, dy, ayVar, jy, &
                   weight*f0vec(iv),  weight*f0vec(iv), iBcPerpPeriodic )
              dfu(iv,iz,:,ix,iDirection) = fyVec(:)  - weight * f0vec(iv) 
           enddo
        enddo
     end do


   enddo ! iDirection

end subroutine evolve_F_advect

!-------------------------------------------------------------------------------------
subroutine evolve_F_accel(dfu, &                                ! input/output
                          Ez2D, fvc0T)                          ! input
!-------------------------------------------------------------------------------------

  use fft_wrap, only   : fft_forward, fft_backward
  use parameters, only : dt, Dvvvv, iunit

  implicit none

  real,    intent(inout)   :: dfu(jv, jz, jy, jx, jDirection)
  real,    intent(in)      :: Ez2D(jz,jy,jx)
  complex, intent(in)      :: fvc0T(jv)

  complex                  :: fvc(jv), z(jv)
  integer                  :: ix, iy, iz, iv, iDirection


  !-- acceleration: df/dt = Ez df/dv

   do iDirection = 1,jDirection
      do ix=1,jx
         do iy=1,jy
            do iz=1,jz   !	begin spatial loop
               fvc(:) = dfu(:,iz,iy,ix,iDirection)
               call fft_forward(jv,fvc,fvc)
               fvc = fvc/float(jv)
               call dealias(jv,fvc)
               z(:)   =  - Dvvvv*pvec(:)**4  - iunit*pvec(:)*Ez2D(iz,iy,ix) + 1.E-09
               fvc(:) = fvc(:)*cexp(z(:)*dt) + &
                    weight * (1-cexp(z(:)*dt))*iunit*pvec(:)*Ez2D(iz,iy,ix)*fvc0T(:)/z(:)
!	       do ip=1,np !  but not ip = 1 since pvec(1) = z(1) = 0 
!	       fvc(ip)=fvc(ip)*cexp(z(ip)*dt) + &
!                     (1-cexp(z(ip)*dt))*iunit*pvec(ip)*Ez2D(ix,iz)*fvc0T(ip)/z(ip)
!              enddo
               call fft_backward(jv,fvc,fvc)
               dfu(:,iz,iy,ix,iDirection)=fvc(:)       
               !   f(iv, iDirection) = dfu(iv, iDirection) + 0.5 * f0vec(iv)
            enddo
         enddo
      enddo
   enddo
   

end subroutine evolve_F_accel

!-------------------------------------------------------------------------------------
subroutine evolve_Uz(dfu, f0vec, &                             ! input
                     uzMat)                                    ! output
!-------------------------------------------------------------------------------------

  implicit none

  real,    intent(in)      :: dfu(jv, jz, jy, jx, jDirection)
  real,    intent(in)      :: f0vec(jv)
  real,    intent(out)     :: uzMat(jz, jy, jx, jDirection)


  real       :: fVsVvec(jv)
  integer    :: ix, iy, iz, iv, iDirection


  !-- integrate over "v" to find longtitudinal velocity Uz -- 
       
  do ix = 1,jx
     do iy = 1, jy
        do iz = 1,jz 
           do iDirection = 1,jDirection
              fVsVvec(:) = dfu(:,iz,iy,ix,iDirection) + weight * f0vec(:)   ! weight or simpy 0.5?
              uzMat(iz,iy,ix,iDirection) = sum( fVsVvec * vVec ) / sum( fVsVvec )
           enddo
        enddo
     enddo
  enddo
   

end subroutine evolve_Uz

!-------------------------------------------------------------------------------------
 
subroutine evolve_U_advect(uxMat, uyMat, &                            ! input/output
                           uzMat)                                     ! input

!-------------------------------------------------------------------------------------

  use fft_wrap, only   : fft_forward, fft_backward

  use parameters, only : ux0, uy0, dt, iWaveFrame, iBcPerpPeriodic, iFourier


  implicit none

  real,    intent(inout)   :: uxMat(jz, jy, jx, jDirection)
  real,    intent(inout)   :: uyMat(jz, jy, jx, jDirection)
  real,    intent(in)      :: uzMat(jz, jy, jx, jDirection)


  complex    :: dummy1Dz(jz)
  complex    :: dummy1Dy(jy)
  complex    :: dummy1Dx(jx)

  real       :: uxVecX(jx)
  real       :: uyVecX(jx)

  real       :: uxVecY(jy)
  real       :: uyVecY(jy)

  real       :: uxVecZ(jz)
  real       :: uyVecZ(jz)
  real       :: uzVecZ(jz)


  integer    :: ix, iy, iz, iv, iDirection

   
  !-- advect U in z:  dUx/dt = -Uz dUx/dz, dUy/dt = -Uz dUy/dz
    
  do iDirection = 1,jDirection
   
     if(iWaveFrame==1)  then
 
        !  advect ux in z since <u_z> approximately = -vphase =  V0, added on June 2, 2009
        do ix=1,jx
           do iy=1,jy

              if( iFourier==1 ) then                   ! FT with linearization in uz

                 dummy1Dz(:) = uxMat(:,iy,ix,iDirection)
                 call fft_forward(jz,dummy1Dz,dummy1Dz)
                 dummy1Dz=dummy1Dz/float(jz)
                 call dealias(jz,dummy1Dz) 
                 dummy1Dz(:) = ExpKz(:) * dummy1Dz(:)
                 call fft_backward(jz,dummy1Dz,dummy1Dz)
                 uxMat(:,iy,ix,iDirection) = dummy1Dz(:)

                 dummy1Dz(:) = uyMat(:,iy,ix,iDirection)
                 call fft_forward(jz,dummy1Dz,dummy1Dz)
                 dummy1Dz=dummy1Dz/float(jz)
                 call dealias(jz,dummy1Dz) 
                 dummy1Dz(:) = ExpKz(:) * dummy1Dz(:)
                 call fft_backward(jz,dummy1Dz,dummy1Dz)
                 uyMat(:,iy,ix,iDirection) = dummy1Dz(:)


              else                                     ! upwind or central differences

                 uzVecZ(:) = uzMat(:,iy,ix,iDirection)
                 uxVecZ(:) = uxMat(:,iy,ix,iDirection)
                 uyVecZ(:) = uyMat(:,iy,ix,iDirection)

                 call SubUpwind_NonCon_Form( uxVecZ, dt, dz, uzVecZ, jz, &
                      0.0, 0.0, 1, -iFourier, dummy1Dz )

                 call SubUpwind_NonCon_Form( uyVecZ, dt, dz, uzVecZ, jz, &
                      0.0, 0.0, 1, -iFourier, dummy1Dz )

                 uxMat(:,iy,ix,iDirection)= uxVecZ(:)
                 uyMat(:,iy,ix,iDirection)= uyVecZ(:)
 
             endif

           end do
        enddo

     endif  ! else need to implement centered difference variant of SubUpwind_NonCon_Form

  end do



  !-- advect U in y:  dUx/dt = -Uy dUx/dy,   dUy/dt = -Uy dUy/dy


  do iDirection = 1,jDirection
   
     if(iWaveFrame==1)  then
 
        do ix=1,jx
           do iz=1,jz
              
              uxVecY(:) = uxMat(iz,:,ix,iDirection)
              uyVecY(:) = uyMat(iz,:,ix,iDirection)

              call SubUpwind_NonCon_Form( uxVecY, dt, dy, uyVecY, jy, &
                   0.0, 0.0, 1, 1, dummy1Dy )

              call NonLinearAdvection( uyVecY, dt, dy, jy, &
                   uy0(iDirection), uy0(iDirection), iBcPerpPeriodic )  

              uxMat(iz,:,ix,iDirection) = uxVecY(:)
              uyMat(iz,:,ix,iDirection) = uyVecY(:)

           end do
        enddo

     endif  ! else need to implement centered difference variant of SubUpwind_NonCon_Form

  end do


  !-- advect U in x:  dUx/dt = -Ux dUx/dx,   dUy/dt = -Ux dUy/dx

  do iDirection = 1,jDirection

     do iy=1,jy 
        do iz=1,jz

           uxVecX(:) = uxMat(iz,iy,:,iDirection)
           uyVecX(:) = uyMat(iz,iy,:,iDirection)

           call SubUpwind_NonCon_Form( uyVecX, dt, dx, uxVecX, jx, &
                   0.0, 0.0, 1, 1, dummy1Dx )

           call NonLinearAdvection( uxVecX, dt, dx, jx, &
                ux0(iDirection), ux0(iDirection), iBcPerpPeriodic )  
 
           uxMat(iz,iy,:,iDirection) = uxVecX(:)
           uyMat(iz,iy,:,iDirection) = uyVecX(:)

        end do
     end do

  end do

end subroutine evolve_U_advect


!-------------------------------------------------------------------------------------
 subroutine evolve_U_accel(uxMat, uyMat, &                     ! input/output
                           Ex2D, Ey2D)                         ! input
!-------------------------------------------------------------------------------------

  use parameters, only : dt

  implicit none

  real,    intent(inout)   :: uxMat(jz, jy, jx, jDirection)
  real,    intent(inout)   :: uyMat(jz, jy, jx, jDirection)

  real,    intent(in)      :: Ex2D(jz,jy,jx)
  real,    intent(in)      :: Ey2D(jz,jy,jx)

  integer                  :: iDirection

  !-- acceleration:  dUx/dt = Ex,  dUy/dt = Ey


  do iDirection = 1,jDirection
   
     uxMat(:,:,:,iDirection) = uxMat(:,:,:,iDirection) + Ex2D * dt
     uyMat(:,:,:,iDirection) = uyMat(:,:,:,iDirection) + Ey2D * dt
         
  enddo
 

end subroutine evolve_U_accel


!-------------------------------------------------------------------------------------
 subroutine evolve_U_filter(uxMat, uyMat, &                     ! input/output
                            ExFilter, EyFilter)                 ! input
!-------------------------------------------------------------------------------------

  use fft_3d

  use parameters, only : ux0, uy0

  implicit none


  real,    intent(inout)   :: uxMat(jz, jy, jx, jDirection)
  real,    intent(inout)   :: uyMat(jz, jy, jx, jDirection)

  real,    intent(in)      :: ExFilter(jx)
  real,    intent(in)      :: EyFilter(jy)

  complex    :: du(jz,jy,jx)

  integer    :: ix, iy, iDirection


  do iDirection = 1,jDirection

     !-- filtering x-velocity

     du = uxMat(:,:,:,iDirection) - ux0(iDirection) 


     do ix=1,jx
        du(:,:,ix) = du(:,:,ix)*uxMod(ix)
     end do

     call fft_forward_x(du)

     do ix=1,jx
        du(:,:,ix) = du(:,:,ix)*ExFilter(ix)
     end do

     call fft_backward_x(du)

     uxMat(:,:,:,iDirection) = du +  ux0(iDirection)


     !-- filtering y-velocity

     du = uyMat(:,:,:,iDirection) - uy0(iDirection) 


     do iy=1,jy
        du(:,iy,:) = du(:,iy,:)*uyMod(iy)
     end do

     call fft_forward_y(du)

     do iy=1,jy
        du(:,iy,:) = du(:,iy,:)*EyFilter(iy)
     end do

     call fft_backward_y(du)

     uyMat(:,:,:,iDirection) = du +  uy0(iDirection)

  enddo


  if (jx .eq. 1) uxMat = 0.
  if (jy .eq. 1) uyMat = 0.

end subroutine evolve_U_filter
 

!-------------------------------------------------------------------------------------

end module evolve
